from mindspore import Tensor, context
import numpy as np

ALL_LOG = []
test_time = 10000
last_cnt = 0
ALL_LOG.append(
    '%-40s\t'
    '%-20s\t'
    '%-20s\t'
    '%-20s\t'
    '%-20s\t'
    '%-20s' % ('name', 'cost_time', 'run_op_time', 'run_op_ratio', 'tuple_time', 'run_op_ratio(no tuple time)')
)
need_profile = False


def timer(func):
    def func_wrapper(*args, **kwargs):
        global last_cnt
        import time
        # warm up
        func(1000)
        if need_profile:
            from mindspore.ops.primitive import run_op_cnt, all_time
            from mindspore.ops.composite.multitype_ops._compile_utils import tuple_time
        else:
            run_op_cnt = 0
            all_time = 0
            tuple_time = 0
        last_cnt = run_op_cnt
        last_time = all_time
        last_tuple_time = tuple_time
        time_start = time.perf_counter()
        result = func(test_time)
        time_end = time.perf_counter()
        time_spend = time_end - time_start
        if need_profile:
            from mindspore.ops.primitive import run_op_cnt, all_time
            from mindspore.ops.composite.multitype_ops._compile_utils import tuple_time
        run_op_time = (all_time - last_time) * 1000
        cost_time = time_spend * 1000
        tuple_pre_time = (tuple_time - last_tuple_time) * 1000
        ALL_LOG.append(
            '%-40s\t'
            '%-20s\t'
            '%-20s\t'
            '%-20s\t'
            '%-20s\t'
            '%-20s' % (func.__name__,
                       str(round(cost_time, 3)) + 'ms',
                       str(round(run_op_time, 3)) + 'ms',
                       str(round(run_op_time * 100 / cost_time, 3)) + '%',
                       str(round(tuple_pre_time, 3)) + 'ms',
                       str(round(run_op_time * 100 / (cost_time - tuple_pre_time), 3)) + '%'))
        last_cnt = run_op_cnt
        return result

    return func_wrapper


@timer
def test_slice_slice(test_time):
    input = Tensor(np.random.rand(64, 12, 4, 4).astype(np.float32))
    for i in range(test_time):
        a = input[1:1 + 1:1, ::]


@timer
def test_ellipsis_slice(test_time):
    input = Tensor(np.random.rand(64, 12, 4, 4).astype(np.float32))
    for i in range(test_time):
        a = input[..., 2:4]


@timer
def test_slice_slice_slice_slice(test_time):
    input = Tensor(np.random.rand(64, 12, 4, 4).astype(np.float32))
    for i in range(test_time):
        a = input[3:15, :, 1:2, 2:4]


@timer
def test_none(test_time):
    input = Tensor(np.random.rand(64, 12, 4, 4).astype(np.float32))
    for i in range(test_time):
        a = input[None]


@timer
def test_int(test_time):
    input = Tensor(np.random.rand(64, 12, 4, 4).astype(np.float32))
    for i in range(test_time):
        a = input[i % 64]


@timer
def test_tensor(test_time):
    input = Tensor(np.random.rand(64, 12, 4, 4).astype(np.float32))
    index = Tensor(np.array([1, 2, 3, 4]).astype(np.int32))
    for i in range(test_time):
        a = input[index]


@timer
def test_list(test_time):
    input = Tensor(np.random.rand(64, 12, 4, 4).astype(np.float32))
    for i in range(test_time):
        a = input[[1, 2, 3, 4]]


@timer
def test_tensor_list_int_slice(test_time):
    input = Tensor(np.random.rand(64, 12, 4, 4).astype(np.float32))
    index = Tensor(np.array([1, 12, 23, 34]).astype(np.int32))
    for i in range(test_time):
        a = input[index, [1, 2, 3, 4], 1, 2:3]


@timer
def test_bool(test_time):
    input = Tensor(np.random.rand(64, 12, 4, 4).astype(np.float32))
    for i in range(test_time):
        a = input[True]


@timer
def test_slice(test_time):
    input = Tensor(np.random.rand(64, 12, 4, 4).astype(np.float32))
    for i in range(test_time):
        a = input[7:17]


if __name__ == '__main__':
    # context.set_context(mode=context.GRAPH_MODE)
    context.set_context(mode=context.PYNATIVE_MODE)
    test_slice_slice()
    test_ellipsis_slice()
    test_slice_slice_slice_slice()
    test_none()
    test_int()
    test_tensor()
    test_list()
    test_tensor_list_int_slice()
    test_bool()
    test_slice()
    for log in ALL_LOG:
        print(log)
